package pool

// Version 3：并发任务池 + context 控制

import (
	"context"
	"sync"
)

type Task func(context.Context) error

type Pool struct {
	workers int
	tasks   []Task
}

func New(workers int) *Pool {
	return &Pool{workers: workers, tasks: make([]Task, 0)}
}

func (p *Pool) Add(task Task) {
	p.tasks = append(p.tasks, task)
}

func (p *Pool) Run(ctx context.Context) error {
	childCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	tasks := make(chan Task)
	errCh := make(chan error, 1)
	var once sync.Once

	worker := func() {
		for {
			select {
			case <-childCtx.Done():
				once.Do(func() { errCh <- childCtx.Err() })
				return
			case task, ok := <-tasks:
				if !ok {
					return
				}
				if err := task(childCtx); err != nil {
					once.Do(func() { errCh <- err })
					cancel()
					return
				}
			}
		}
	}

	var wg sync.WaitGroup
	for i := 0; i < p.workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			worker()
		}()
	}

	go func() {
		defer close(tasks)
		for _, task := range p.tasks {
			select {
			case <-childCtx.Done():
				return
			case tasks <- task:
			}
		}
	}()

	wg.Wait()

	select {
	case err := <-errCh:
		return err
	default:
		return nil
	}
}
